#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

import UpdateLMSG_Module
from pymongo import MongoClient, IndexModel, ASCENDING, DESCENDING
import re
import json
import sys

HOST = "10.188.188.22"
PORT = 27017
DB = "eLMSG"
COLL = "eLMSG_20180612"
nameset = set(["unC", "unO", "unF", "unP", "UnC", "UnO", "UnF", "UnP"])


# 连接MongoDB
def InitMongoDB(MONGODB_HOST, MONGODB_PORT, MONGODB_DB_NAME, MONGODB_COLLECTION_NAME):
    client = MongoClient(MONGODB_HOST, MONGODB_PORT)
    db = client[MONGODB_DB_NAME]
    collection = db[MONGODB_COLLECTION_NAME]
    return client, collection


# 关闭MongoDB
def CloseMongoDB(MongoDB_Connection):
    MongoDB_Connection.close()


# 列出当前库中所有非Strain的唯一名称，返回set
def LoadCurrentNames(MONGODB_HOST, MONGODB_PORT, MONGODB_DB_NAME, MONGODB_COLLECTION_NAME):
    names = set()

    client, collection = InitMongoDB(MONGODB_HOST, MONGODB_PORT, MONGODB_DB_NAME, MONGODB_COLLECTION_NAME)

    for result in collection.find({"Rank": {"$ne": "Strain"}}, {"Name": 1}):
        names.add(result["Name"])

    CloseMongoDB(client)
    return names


# 获取一条Lineage str的parent name
def GetParentName(lineage_str: str):
    nameset = set(["unC", "unO", "unF", "unP", "UnC", "UnO", "UnF", "UnP"])
    items = re.search("\(.*\)", lineage_str)
    if items != None:
        lineage = re.sub("\.\(.*\)", "", lineage_str).split(".")
        for i in range(len(lineage) - 1, -1, -1):
            if lineage[i] in nameset:
                lineage.pop(i)
        if len(lineage) > 1:
            return lineage[-1]
    else:
        lineage = lineage_str.split(".")
        for i in range(len(lineage) - 1, -1, -1):
            if lineage[i] in nameset:
                lineage.pop(i)
        if len(lineage) > 2:
            return lineage[-2]
    return None


# 清除空字段
def del_empty_feild(_dict: dict, parent=None, key_in_parent=None):
    for k in list(_dict.keys()):
        v = _dict[k]
        if isinstance(v, str):
            if v is None or not len(v.strip()):
                del _dict[k]
        elif isinstance(v, list):
            if not len(v):
                del _dict[k]
            else:
                for i in v:
                    if isinstance(i, dict):
                        del_empty_feild(i)
        elif isinstance(v, dict):
            del_empty_feild(v, _dict, k)
        elif v is None:
            del _dict[k]

    if not len(_dict) and parent is not None and key_in_parent is not None:
        del parent[key_in_parent]


def Clean_eLMSG():
    client_eLMSG, collection_eLMSG = InitMongoDB(HOST, PORT, DB, COLL)

    for result in collection_eLMSG.find():
        del_empty_feild(result)
        collection_eLMSG.replace_one({"_id": result["_id"]}, result)

    CloseMongoDB(client_eLMSG)


# 补充LMSG_ID
def FullfillLMSG_ID(startID):
    client_eLMSG, collection_eLMSG = InitMongoDB(HOST, PORT, DB, COLL)

    for result in collection_eLMSG.find({"LMSG_ID": {"$exists": 0}}, {"_id": 1}):
        startID += 1
        lmsg_id = "MSG" + str(startID).zfill(6)
        collection_eLMSG.update_one({"_id": result["_id"]}, {"$set": {"LMSG_ID": lmsg_id}})

    CloseMongoDB(client_eLMSG)


# 补充Parent LMSG ID
def FullfillParentLMSGId():
    ht = {}
    client_eLMSG, collection_eLMSG = InitMongoDB(HOST, PORT, DB, COLL)
    for result in collection_eLMSG.find({"Rank": {"$ne": "Strain"}}, {"Name": 1, "LMSG_ID": "1", "_id": 0}):
        ht[result["Name"]] = result["LMSG_ID"]
    for result in collection_eLMSG.find({"Parent_LMSG_ID": {"$exists": 0}}, {"Lineage": 1}):
        items = re.search("\(.*\)", result["Lineage"])
        if items != None:
            lineage = re.sub("\.\(.*\)", "", result["Lineage"]).split(".")
            for i in range(len(lineage) - 1, -1, -1):
                if lineage[i] in nameset:
                    lineage.pop(i)
            if len(lineage) > 1:
                collection_eLMSG.update_one({"_id": result["_id"]}, {"$set": {"Parent_LMSG_ID": ht[lineage[-1]]}})
        else:
            lineage = result["Lineage"].split(".")
            for i in range(len(lineage) - 1, -1, -1):
                if lineage[i] in nameset:
                    lineage.pop(i)
            if len(lineage) > 2:
                collection_eLMSG.update_one({"_id": result["_id"]}, {"$set": {"Parent_LMSG_ID": ht[lineage[-2]]}})
    CloseMongoDB(client_eLMSG)


# 生成结构化的Reference
def GenerateStructuredReference():
    client_eLMSG, collection_eLMSG = InitMongoDB(HOST, PORT, DB, COLL)

    for result in collection_eLMSG.find({"References": {"$exists": 1}, "StructuredReference": {"$exists": 0}},
                                        {"References": 1}):
        reference = []
        for line in result["References"]:
            reference.append(GetReference(line))
        collection_eLMSG.update_one({"_id": result["_id"]}, {"$set": {"StructuredReference": reference}})
    CloseMongoDB(client_eLMSG)


def GetReference(line):
    import re
    reference = {}
    regex = "\(((PM|DO).*)\)"
    prog = re.compile(regex)
    result = prog.search(line)
    if result:
        ids = result.group(1)
        for id in re.split(r'[\s ]', ids):
            if id.startswith("PMID"):
                reference["PMID"] = id[5:]
            elif id.startswith("DOI"):
                doi = id[4:]
                doi = re.sub("<.*?>", "", doi)
                reference["DOI"] = doi
            elif id.startswith("PMCID"):
                reference["PMCID"] = id[6:]
        txt = line.replace("(" + ids + ")", "", -1).strip()
        if txt != "":
            reference["txt"] = txt
    else:
        txt = line.strip()
        if txt != "":
            reference["txt"] = txt
    return reference


# 结构化Genbank_16SrRNA
def Structure_Genbank_16SrRNA():
    client_eLMSG, collection_eLMSG = InitMongoDB(HOST, PORT, DB, COLL)
    for result in collection_eLMSG.find({"16SrRNA": {"$exists": 1}}, {"16SrRNA": 1}):
        collection_eLMSG.update_one({"_id": result["_id"]},
                                    {"$set": {"Genbank_16SrRNA": result["16SrRNA"]}, "$unset": {"16SrRNA": 1}})
    CloseMongoDB(client_eLMSG)


# 结构化JGI_IMG
def Structure_JGI_IMG():
    client_eLMSG, collection_eLMSG = InitMongoDB(HOST, PORT, DB, COLL)
    for result in collection_eLMSG.find({"JGI IMG": {"$exists": 1}}, {"JGI IMG": 1}):
        collection_eLMSG.update_one({"_id": result["_id"]},
                                    {"$set": {"JGI_IMG": result["JGI IMG"]}, "$unset": {"JGI IMG": 1}})
    CloseMongoDB(client_eLMSG)


# 结构化NCBI_Assembly
def Structure_NCBI_Assembly():
    client_eLMSG, collection_eLMSG = InitMongoDB(HOST, PORT, DB, COLL)
    for result in collection_eLMSG.find({"NCBI GenBank": {"$exists": 1}}, {"NCBI GenBank": 1}):
        collection_eLMSG.update_one({"_id": result["_id"]},
                                    {"$set": {"NCBI_Assembly": result["NCBI GenBank"]}, "$unset": {"NCBI GenBank": 1}})
    CloseMongoDB(client_eLMSG)


# 更改不标准的Type字段结构
def ChangeTypeFieldType():
    client_eLMSG, collection_eLMSG = InitMongoDB(HOST, PORT, DB, COLL)

    for result in collection_eLMSG.find({"Type": {"$exists": 1}}, {"Type": 1}):
        if isinstance(result["Type"], str):
            type = result["Type"]
            id = result["_id"]
            collection_eLMSG.update_one({"_id": id}, {"$set": {"Type": [type]}})

    CloseMongoDB(client_eLMSG)


# 更改不标准的Phenotype.Physiology.nutrition_type字段结构
def ChangeNutritionTypeFieldType():
    client_eLMSG, collection_eLMSG = InitMongoDB(HOST, PORT, DB, COLL)

    for result in collection_eLMSG.find({"Phenotype.Physiology.nutrition_type": {"$exists": 1}},
                                        {"Phenotype.Physiology.nutrition_type": 1}):
        if isinstance(result["Phenotype"]["Physiology"]["nutrition_type"], str):
            type = result["Phenotype"]["Physiology"]["nutrition_type"]
            id = result["_id"]
            collection_eLMSG.update_one({"_id": id}, {"$set": {"Phenotype.Physiology.nutrition_type": [type]}})

    CloseMongoDB(client_eLMSG)


# 更改Phenotype.Biochemistry.fattyacids字段结构
def StructureFattyAcids():
    client_eLMSG, collection_eLMSG = InitMongoDB(HOST, PORT, DB, COLL)
    for result in collection_eLMSG.find({"Phenotype.Biochemistry.fattyacids": {"$exists": 1}},
                                        {"Phenotype.Biochemistry.fattyacids": 1}):
        if isinstance(result["Phenotype"]["Biochemistry"]["fattyacids"], str):
            fattyacids = trans_fattyacids(result["Phenotype"]["Biochemistry"]["fattyacids"])
            collection_eLMSG.update_one({"_id": result["_id"]},
                                        {"$set": {"Phenotype.Biochemistry.fattyacids": fattyacids}})
    CloseMongoDB(client_eLMSG)


def trans_fattyacids(fattyacids_str):
    import json
    fattyacids = []
    for k, v in json.loads(fattyacids_str).items():
        fattyacids.append({"name": k, "account": str(v)})
    return fattyacids


# 增加基因组注释数据
def AddFunctionAnnotation(filepath):
    client_eLMSG, collection_eLMSG = InitMongoDB(HOST, PORT, DB, COLL)

    infile = open(filepath, "r")
    infile.readline()
    for line in infile:
        line = line.rstrip()
        splitline = line.split("\t")
        if splitline[0].startswith("GCA_"):
            result = collection_eLMSG.find_one({"Rank": "Strain", "NCBI_Assembly.id": splitline[0]}, {"_id": 1})
            if result != None:
                gcaId = splitline[0].replace(".", "_")
                anno = {"Swissprot": int(splitline[4]), "TIGRFAMs": int(splitline[5]), "KEGG": int(splitline[2]),
                        "COGs": int(splitline[1]),
                        "Pfam": int(splitline[3]), "GO": int(splitline[6]), "EC": int(splitline[7]),
                        "MetaCyc": int(splitline[8])}
                collection_eLMSG.update_one({"_id": result["_id"]}, {"$set": {"FunctionAnnotation." + gcaId: anno}})
        else:
            result = collection_eLMSG.find_one({"Rank": "Strain", "JGI_IMG.id": splitline[0]}, {"_id": 1})
            if result != None:
                jgiId = splitline[0].replace(".", "_")
                anno = {"Swissprot": int(splitline[4]), "TIGRFAMs": int(splitline[5]), "KEGG": int(splitline[2]),
                        "COGs": int(splitline[1]),
                        "Pfam": int(splitline[3]), "GO": int(splitline[6]), "EC": int(splitline[7]),
                        "MetaCyc": int(splitline[8])}
                collection_eLMSG.update_one({"_id": result["_id"]}, {"$set": {"FunctionAnnotation." + jgiId: anno}})
    infile.close()
    CloseMongoDB(client_eLMSG)


# 增加JGI_IMG元数据
def FillIMGMeta(filepath):
    client_eLMSG, collection_eLMSG = InitMongoDB(HOST, PORT, DB, COLL)

    infile = open(filepath, "r")
    infile.readline()
    # count=2
    for line in infile:
        line = line.rstrip("\n")
        splitline = line.split("\t")
        # print(str(count)+":"+str(len(splitline)))
        # count+=1
        result = collection_eLMSG.find({"JGI_IMG.id": splitline[0].strip()[4:]}, {"_id": 1})
        if result != []:
            meta = {}
            if splitline[1].strip() != "":
                meta["JGI_IMG.$.project"] = splitline[1].strip()
            if splitline[2].strip() != "":
                meta["JGI_IMG.$.principal_investigator"] = splitline[2].strip()
            if splitline[3].strip() != "":
                meta["JGI_IMG.$.email"] = splitline[3].strip()
            if splitline[4].strip() != "":
                meta["JGI_IMG.$.sequencing_center"] = splitline[4].strip()
            if splitline[5].strip() != "":
                meta["JGI_IMG.$.submitter_date"] = splitline[5].strip()
            if splitline[6].strip() != "":
                meta["JGI_IMG.$.update_date"] = splitline[6].strip()
            if splitline[7].strip() != "":
                meta["JGI_IMG.$.submission_type"] = splitline[7].strip()
            if splitline[8].strip() != "":
                meta["JGI_IMG.$.sequencing_status"] = splitline[8].strip()
            if splitline[9].strip() != "":
                meta["JGI_IMG.$.seq_status"] = splitline[9].strip()
            if splitline[10].strip() != "":
                xref = []
                ids = splitline[10].strip().split(";")
                for id in ids:
                    if id.strip().startswith("PRJ"):
                        xref.append({"db": "bioproject", "id": id.strip()})
                    elif id.strip().startswith("SAM"):
                        xref.append({"db": "biosample", "id": id.strip()})
                if len(xref) != 0:
                    meta["JGI_IMG.$.xref"] = xref
            collection_eLMSG.update_many({"JGI_IMG": {"$elemMatch": {"id": splitline[0].strip()[4:]}}},
                                         {"$set": meta})
    infile.close()

    CloseMongoDB(client_eLMSG)


# 测试非Strain新增数据的完整性
def TestNoStrainNovel(fileList: list) -> bool:
    passed = True
    tmpNames = set()

    client, collection = InitMongoDB(HOST, PORT, DB, COLL)

    for file in fileList:
        infile = open(file, "r", encoding='UTF-8')
        for line in infile:
            line = line.strip()
            json_data = json.loads(line, encoding='utf-8')
            if "Name" in json_data:
                result = collection.find_one({"Name": json_data["Name"]})
                if result != None:
                    print("Some record named \"" + json_data["Name"] + "\" already exists.")
                    passed = False
                else:
                    if "Lineage" in json_data:
                        parentName = GetParentName(json_data["Lineage"])
                        if parentName == None:
                            print("Cannot get the parent name for \"" + json_data["Name"] + "\".")
                            passed = False
                        else:
                            result2 = collection.find_one({"Name": parentName})
                            if result2 == None:
                                if parentName not in tmpNames:
                                    print("Cannot find the parent record for \"" + json_data[
                                        "Name"] + "\". The parent name is \"" + parentName + "\".")
                                    passed = False
                                else:
                                    tmpNames.add(json_data["Name"])
                            else:
                                tmpNames.add(json_data["Name"])
                    else:
                        print("Some record has no lineage in \"" + file + "\".")
                        passed = False
            else:
                print("Some record has no name in \"" + file + "\".")
                passed = False
        infile.close()
    CloseMongoDB(client)
    if passed:
        print("Check no strain novel files passed.")
    else:
        print("Invalid no strain novel files, check failed.")
    return passed


# 测试Strain新增数据的完整性
def TestStrainNovel(fileList: list) -> bool:
    passed = True
    tmpNames = set()

    client, collection = InitMongoDB(HOST, PORT, DB, COLL)

    for file in fileList:
        infile = open(file, "r", encoding='UTF-8')
        for line in infile:
            line = line.strip()
            json_data = json.loads(line, encoding='utf-8')
            if "Name" in json_data:
                if "Lineage" in json_data:
                    parentName = GetParentName(json_data["Lineage"])
                    if parentName == None:
                        print("Cannot get the parent name for \"" + json_data["Name"] + "\".")
                        passed = False
                    else:
                        result2 = collection.find_one({"Name": parentName})
                        if result2 == None:
                            print("Cannot find the parent record for \"" + json_data[
                                "Name"] + "\". The parent name is \"" + parentName + "\".")
                            passed = False
                else:
                    print("Some record has no lineage in \"" + file + "\".")
                    passed = False
            else:
                print("Some record has no name in \"" + file + "\".")
                passed = False
        infile.close()
    CloseMongoDB(client)
    if passed:
        print("Check no strain novel files passed.")
    else:
        print("Invalid no strain novel files, check failed.")
    return passed


# 获取当前库最大LMSG ID
def GetMaxLMSG_ID() -> int:
    client, collection = InitMongoDB(HOST, PORT, DB, COLL)
    for result in collection.find().sort("LMSG_ID", DESCENDING):
        max_id = int(result["LMSG_ID"][3:])
        break
    CloseMongoDB(client)
    return max_id


# 导入非Strain新增数据
def InsertNoStrainNovel(fileList: str):
    client, collection = InitMongoDB(HOST, PORT, DB, COLL)
    for file in fileList:
        infile = open(file, "r", encoding='UTF-8')
        for line in infile:
            line = line.strip()
            json_data = json.loads(line, encoding='utf-8')
            collection.find_one_and_replace({"Name": json_data["Name"]}, json_data, upsert=True)
    FullfillLMSG_ID(GetMaxLMSG_ID())
    FullfillParentLMSGId()
    CloseMongoDB(client)


# 导入Strain新增数据
def InsertStrainNovel(fileList: str):
    client, collection = InitMongoDB(HOST, PORT, DB, COLL)
    for file in fileList:
        infile = open(file, "r", encoding='UTF-8')
        for line in infile:
            line = line.strip()
            json_data = json.loads(line, encoding='utf-8')
            collection.find_one_and_replace({"Name": json_data["Name"]}, json_data, upsert=True)
    FullfillLMSG_ID(GetMaxLMSG_ID())
    FullfillParentLMSGId()

    CloseMongoDB(client)


# 导入新增数据后标准化数据结构和补充部分内容
def DatabaseNormalization(IMG_genome_metadata_filepath: str, function_annotation_filepath: str):
    print("Cleaning eLMSG and remove null nodes ...")
    Clean_eLMSG()

    print("Generating structured references ...")
    GenerateStructuredReference()

    print("Generating structured Genbank_16SrRNA ...")
    Structure_Genbank_16SrRNA()

    print("Generating structured JGI_IMG ...")
    Structure_JGI_IMG()

    print("Generating structured NCBI_Assembly ...")
    Structure_NCBI_Assembly()

    print("Generating normalized type ...")
    ChangeTypeFieldType()

    print("Generating normalized nutrition type ...")
    ChangeNutritionTypeFieldType()

    print("Generating normalized fatty acids ...")
    StructureFattyAcids()

    print("Filling JGI_IMG metadata ...")
    FillIMGMeta(IMG_genome_metadata_filepath)

    print("Adding funtional annotations ...")
    AddFunctionAnnotation(function_annotation_filepath)


def create_structured_lineage(lmsg_id: str):
    client, collection = InitMongoDB(HOST, PORT, DB, COLL)

    def d(me, out):
        out.add(me["Name"])

        if "Parent_LMSG_ID" not in me:
            return

        parent = collection.find_one({"LMSG_ID": me["Parent_LMSG_ID"]},
                                     {"Name": 1, "Parent_LMSG_ID": 1})
        if parent is None:
            return

        d(parent, out)
    print(lmsg_id)
    it = collection.find_one({"LMSG_ID": lmsg_id}, {"Parent_LMSG_ID": 1, "Name": 1})
    structured_lineage = set()
    d(it, structured_lineage)

    structured_lineage = list(structured_lineage)
    collection.update_one({"_id": it["_id"]}, {"$set": {"StructuredLineage": structured_lineage}})

    CloseMongoDB(client)


# 更新非Strain更新数据
def UpdateNoStrainUpdates(fileList: str):
    print("Update no strain updates ...")
    null = None
    client, collection = InitMongoDB(HOST, PORT, DB, COLL)
    for file in fileList:
        infile = open(file, "r", encoding='UTF-8')
        for line in infile:
            if line.startswith("New"):
                splitline = line.split("\t")
                lmsg_id = splitline[1][8:].strip()
                name = splitline[2][5:].strip()
                update_key = splitline[3][4:].strip()
                update_value = splitline[4][6:].strip()
                if update_value.startswith("[") or update_value.startswith("{"):
                    update_value = eval(update_value)
                collection.update_one({"LMSG_ID": lmsg_id}, {"$set": {update_key: update_value}})
                if update_key == "References":
                    reference = []
                    for strr in update_value:
                        reference.append(GetReference(strr))
                    collection.update_one({"LMSG_ID": lmsg_id}, {"$set": {"StructuredReference": reference}})
                elif update_key == "Lineage":
                    items = re.search("\(.*\)", update_value)
                    if items != None:
                        lineage = re.sub("\.\(.*\)", "", update_value).split(".")
                        for i in range(len(lineage) - 1, -1, -1):
                            if lineage[i] in nameset:
                                lineage.pop(i)
                        if len(lineage) > 1:
                            parent_record = collection.find_one({"Name": lineage[-1], "Rank": {"$ne": "Strain"}},
                                                                {"LMSG_ID": 1})
                            if parent_record != None:
                                collection.update_one({"LMSG_ID": lmsg_id},
                                                      {"$set": {"Parent_LMSG_ID": parent_record["LMSG_ID"]}})
                                create_structured_lineage(lmsg_id)
                    else:
                        lineage = update_value.split(".")
                        for i in range(len(lineage) - 1, -1, -1):
                            if lineage[i] in nameset:
                                lineage.pop(i)
                        if len(lineage) > 2:
                            parent_record = collection.find_one({"Name": lineage[-2], "Rank": {"$ne": "Strain"}},
                                                                {"LMSG_ID": 1})
                            if parent_record != None:
                                collection.update_one({"LMSG_ID": lmsg_id},
                                                      {"$set": {"Parent_LMSG_ID": parent_record["LMSG_ID"]}})
                                create_structured_lineage(lmsg_id)

    CloseMongoDB(client)


if __name__ == '__main__':
    noStrainNovelList = ["./0823/HR.novel.json",
                         "./0823/Species.novel.json"]
    strainNovelList = ["./0823/Strains_novel.json"]
    if TestNoStrainNovel(noStrainNovelList):
        InsertNoStrainNovel(noStrainNovelList)
        if TestStrainNovel(strainNovelList):
            InsertStrainNovel(strainNovelList)
            DatabaseNormalization("./0823/IMG_genome_metadata.txt",
                                  "./0823/annotation_new_stat")

            UpdateLMSG_Module.Yangxiao_pipeline(sys.argv[0], HOST, DB, COLL)
    UpdateNoStrainUpdates(
        ["./0823/HR内容修改_0828.txt", "./0823/species_修改内容_0828.txt"])
